package main

import (
    "fmt"
    "log"
    "net/http"

    "github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
}

type client struct {
    conn *websocket.Conn
    send chan []byte
}

type hub struct {
    clients    map[*client]bool
    broadcast  chan []byte
    register   chan *client
    unregister chan *client
}

func (h *hub) run() {
    for {
        select {
        case client := <-h.register:
            h.clients[client] = true
        case client := <-h.unregister:
            if _, ok := h.clients[client]; ok {
                delete(h.clients, client)
                close(client.send)
            }
        case message := <-h.broadcast:
            for client := range h.clients {
                select {
                case client.send <- message:
                default:
                    close(client.send)
                    delete(h.clients, client)
                }
            }
        }
    }
}

func (c *client) read(h *hub) {
    defer func() {
        h.unregister <- c
        c.conn.Close()
    }()
    for {
        _, message, err := c.conn.ReadMessage()
        if err != nil {
            break
        }
        h.broadcast <- message
    }
}

func (c *client) write() {
    defer c.conn.Close()
    for message := range c.send {
        err := c.conn.WriteMessage(websocket.TextMessage, message)
        if err != nil {
            break
        }
    }
}

func handleWebSocket(h *hub, w http.ResponseWriter, r *http.Request) {
    // 将 HTTP 连接升级为 WebSocket 连接
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Println(err)
        return
    }

    // 创建新的客户端
    client := &client{conn: conn, send: make(chan []byte, 256)}

    // 注册新的客户端
    h.register <- client

    // 启动客户端读写循环
    go client.write()
    client.read(h)
}

func main() {
    h := &hub{
        clients:    make(map[*client]bool),
        broadcast:  make(chan []byte),
        register:   make(chan *client),
        unregister: make(chan *client),
    }

    go h.run()

    http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
        handleWebSocket(h, w, r)
    })

    fmt.Println("WebSocket server started")
    log.Fatal(http.ListenAndServe(":8080", nil))
}
